import os
from typing import Iterable

import numpy as np
import torch
from tqdm import tqdm
from transformers import PreTrainedModel

from zkl_rome import ComputeCCallback, ComputeCHparams, ComputeVDeltaHparams, load_or_compute_c_inv
from .compute_wilke import ComputeWilkeCallback, ComputeWilkeMetrics


def load_or_compute_modules_c_inv(*,
    model: PreTrainedModel,
    modules: Iterable[torch.nn.Module],
    compute_c_samples: Iterable[np.ndarray] | None = None,
    compute_c_hparams: ComputeCHparams | None = None,
    compute_c_callback: ComputeCCallback | None = None,
    cache_modules_c_inv_file_path: Iterable[os.PathLike | str] | None = None,
) -> tuple[torch.Tensor | None, ...]:
    modules = tuple(modules)
    cache_modules_c_inv_file_path = tuple(cache_modules_c_inv_file_path) \
        if cache_modules_c_inv_file_path is not None else [None] * len(modules)
    return tuple(
        load_or_compute_c_inv(
            model=model,
            module=module,
            compute_c_samples=compute_c_samples,
            compute_c_hparams=compute_c_hparams,
            compute_c_callback=compute_c_callback,
            cache_c_inv_file_path=cache_c_inv_file_path)
        for module, cache_c_inv_file_path in zip(modules, cache_modules_c_inv_file_path))


class TqdmComputeWilkeCallback(ComputeWilkeCallback):
    def __init__(self):
        self.progressbar: tqdm | None = None

    def on_start(self, modules_num: int, hparams: ComputeVDeltaHparams):
        self.progressbar = tqdm(total=modules_num, desc="Computing wilke")

    def on_traverse(self, module_index: int, metrics: ComputeWilkeMetrics) -> bool:
        self.progressbar.update(1)
        return True

    def on_stop(self):
        self.progressbar.close()
